#!/usr/bin/env python3
"""
Main driver script for reproducing the work from:
"Dynamic Analysis of Composite Laminated Plates with a Functionally Graded Graphene-Reinforced Layer"
by Wang et al. (2015)

This software implements:
1. Modified Halpin-Tsai model for graphene-reinforced composites
2. Functionally graded material distribution patterns
3. First-order shear deformation theory (FSDT)
4. Differential quadrature method for vibration analysis
5. Various boundary conditions and laminate configurations
"""

import numpy as np
import matplotlib.pyplot as plt
import argparse
import sys
from pathlib import Path

# Import our modules
from fg_grc_plate import (GrapheneReinforcedComposite, MaterialProperties, 
                         GrapheneGeometry, DistributionPattern)
from laminated_plate import LaminatedPlate, Layer
from vibration_analysis import PlateVibrationSolver, BoundaryCondition
from validation_examples import ValidationStudy, create_validation_plots

def create_material_properties():
    """Create material properties from Table 1 of the paper."""
    epoxy = MaterialProperties(E=3.0, nu=0.34, rho=1200)
    graphene = MaterialProperties(E=1010, nu=0.186, rho=1060)
    gpl_geometry = GrapheneGeometry(
        length=2.5e-6,    # 2.5 μm
        width=1.5e-6,     # 1.5 μm
        thickness=1.5e-9  # 1.5 nm
    )
    return epoxy, graphene, gpl_geometry

def parametric_study(output_dir: Path = Path(".")):
    """
    Comprehensive parametric study reproducing key results from the paper.
    """
    print("Running Parametric Study...")
    print("=" * 50)
    
    epoxy, graphene, gpl_geometry = create_material_properties()
    grc = GrapheneReinforcedComposite(epoxy, graphene, gpl_geometry)
    
    # Study parameters
    plate_size = 0.1  # 10 cm square plate
    thickness = 2.0e-3  # 2 mm total thickness
    
    # 1. Effect of GPL volume fraction
    print("\n1. Effect of GPL Volume Fraction")
    volume_fractions = np.linspace(0, 0.03, 16)  # 0 to 3%
    freq_vs_vf = []
    
    for V_G in volume_fractions:
        layers = [Layer(thickness=thickness, material=epoxy, is_grc=True,
                       grc_material=grc, grc_pattern=DistributionPattern.UD,
                       grc_volume_fraction=V_G)]
        
        plate = LaminatedPlate(layers, length=plate_size, width=plate_size)
        solver = PlateVibrationSolver(plate, n_points_x=9, n_points_y=9)
        
        bc_ss = (BoundaryCondition.SIMPLY_SUPPORTED, BoundaryCondition.SIMPLY_SUPPORTED)
        freqs, _ = solver.solve_eigenvalue_problem(bc_ss, bc_ss, n_modes=1)
        freq_vs_vf.append(freqs[0])
    
    # 2. Effect of distribution patterns
    print("\n2. Effect of Distribution Patterns")
    patterns = {
        'UD': DistributionPattern.UD,
        'FG-V': DistributionPattern.FG_V, 
        'FG-X': DistributionPattern.FG_X,
        'FG-O': DistributionPattern.FG_O
    }
    
    pattern_freqs = {}
    V_G_fixed = 0.01  # 1% volume fraction
    
    for pattern_name, pattern in patterns.items():
        layers = [Layer(thickness=thickness, material=epoxy, is_grc=True,
                       grc_material=grc, grc_pattern=pattern,
                       grc_volume_fraction=V_G_fixed)]
        
        plate = LaminatedPlate(layers, length=plate_size, width=plate_size)
        solver = PlateVibrationSolver(plate, n_points_x=9, n_points_y=9)
        
        bc_ss = (BoundaryCondition.SIMPLY_SUPPORTED, BoundaryCondition.SIMPLY_SUPPORTED)
        freqs, _ = solver.solve_eigenvalue_problem(bc_ss, bc_ss, n_modes=5)
        pattern_freqs[pattern_name] = freqs
    
    # 3. Effect of aspect ratio
    print("\n3. Effect of Aspect Ratio")
    aspect_ratios = np.linspace(0.5, 2.0, 11)
    aspect_freqs = []
    
    for ar in aspect_ratios:
        width = plate_size / np.sqrt(ar)
        length = plate_size * np.sqrt(ar)
        
        layers = [Layer(thickness=thickness, material=epoxy, is_grc=True,
                       grc_material=grc, grc_pattern=DistributionPattern.FG_X,
                       grc_volume_fraction=V_G_fixed)]
        
        plate = LaminatedPlate(layers, length=length, width=width)
        solver = PlateVibrationSolver(plate, n_points_x=9, n_points_y=9)
        
        bc_ss = (BoundaryCondition.SIMPLY_SUPPORTED, BoundaryCondition.SIMPLY_SUPPORTED)
        freqs, _ = solver.solve_eigenvalue_problem(bc_ss, bc_ss, n_modes=1)
        aspect_freqs.append(freqs[0])
    
    # 4. Effect of boundary conditions
    print("\n4. Effect of Boundary Conditions")
    boundary_conditions = {
        'SSSS': ((BoundaryCondition.SIMPLY_SUPPORTED, BoundaryCondition.SIMPLY_SUPPORTED),
                 (BoundaryCondition.SIMPLY_SUPPORTED, BoundaryCondition.SIMPLY_SUPPORTED)),
        'CCCC': ((BoundaryCondition.CLAMPED, BoundaryCondition.CLAMPED),
                 (BoundaryCondition.CLAMPED, BoundaryCondition.CLAMPED)),
    }
    
    bc_freqs = {}
    layers = [Layer(thickness=thickness, material=epoxy, is_grc=True,
                   grc_material=grc, grc_pattern=DistributionPattern.FG_X,
                   grc_volume_fraction=V_G_fixed)]
    
    plate = LaminatedPlate(layers, length=plate_size, width=plate_size)
    solver = PlateVibrationSolver(plate, n_points_x=9, n_points_y=9)
    
    for bc_name, (bc_x, bc_y) in boundary_conditions.items():
        freqs, _ = solver.solve_eigenvalue_problem(bc_x, bc_y, n_modes=3)
        bc_freqs[bc_name] = freqs
    
    # Create comprehensive plots
    fig = plt.figure(figsize=(16, 12))
    
    # Plot 1: GPL volume fraction effect
    ax1 = plt.subplot(2, 3, 1)
    enhancement = np.array(freq_vs_vf) / freq_vs_vf[0]
    plt.plot(volume_fractions * 100, enhancement, 'b-', linewidth=2, marker='o')
    plt.xlabel('GPL Volume Fraction (%)')
    plt.ylabel('Frequency Enhancement Factor')
    plt.title('Effect of GPL Volume Fraction')
    plt.grid(True, alpha=0.3)
    
    # Plot 2: Distribution patterns (Mode 1)
    ax2 = plt.subplot(2, 3, 2)
    pattern_names = list(pattern_freqs.keys())
    mode1_freqs = [pattern_freqs[p][0] for p in pattern_names]
    colors = ['blue', 'red', 'green', 'orange']
    bars = plt.bar(pattern_names, mode1_freqs, color=colors, alpha=0.7)
    plt.ylabel('Frequency (rad/s)')
    plt.title('Effect of GPL Distribution (Mode 1)')
    plt.grid(True, alpha=0.3, axis='y')
    
    # Add frequency values on bars
    for bar, freq in zip(bars, mode1_freqs):
        plt.text(bar.get_x() + bar.get_width()/2, bar.get_height() + freq*0.01,
                f'{freq:.1f}', ha='center', va='bottom', fontsize=9)
    
    # Plot 3: Distribution patterns (Multiple modes)
    ax3 = plt.subplot(2, 3, 3)
    x = np.arange(5)  # 5 modes
    width = 0.2
    for i, (pattern_name, freqs) in enumerate(pattern_freqs.items()):
        plt.bar(x + i*width, freqs, width, label=pattern_name, color=colors[i], alpha=0.7)
    plt.xlabel('Mode Number')
    plt.ylabel('Frequency (rad/s)')
    plt.title('Natural Frequencies vs Mode Number')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.xticks(x + width*1.5, [f'{i+1}' for i in range(5)])
    
    # Plot 4: Aspect ratio effect
    ax4 = plt.subplot(2, 3, 4)
    plt.plot(aspect_ratios, aspect_freqs, 'g-', linewidth=2, marker='s')
    plt.xlabel('Aspect Ratio (a/b)')
    plt.ylabel('Frequency (rad/s)')
    plt.title('Effect of Aspect Ratio')
    plt.grid(True, alpha=0.3)
    
    # Plot 5: Boundary conditions
    ax5 = plt.subplot(2, 3, 5)
    bc_names = list(bc_freqs.keys())
    mode1_bc = [bc_freqs[bc][0] for bc in bc_names]
    plt.bar(bc_names, mode1_bc, color=['cyan', 'magenta'], alpha=0.7)
    plt.ylabel('Frequency (rad/s)')
    plt.title('Effect of Boundary Conditions')
    plt.grid(True, alpha=0.3, axis='y')
    
    # Plot 6: Material property comparison
    ax6 = plt.subplot(2, 3, 6)
    materials = ['Epoxy', 'GRC (1%)', 'GRC (2%)', 'GRC (3%)']
    test_vfs = [0, 0.01, 0.02, 0.03]
    moduli = [grc.effective_modulus(vf) for vf in test_vfs]
    
    plt.bar(materials, moduli, color=['gray', 'lightblue', 'blue', 'darkblue'], alpha=0.7)
    plt.ylabel('Effective Modulus (GPa)')
    plt.title('Effective Material Properties')
    plt.xticks(rotation=45)
    plt.grid(True, alpha=0.3, axis='y')
    
    plt.tight_layout()
    plt.savefig(output_dir / 'parametric_study_results.png', dpi=300, bbox_inches='tight')
    plt.show()
    
    # Print summary results
    print(f"\nSummary Results:")
    print(f"Baseline frequency (pure epoxy): {freq_vs_vf[0]:.2f} rad/s")
    print(f"Enhanced frequency (1% GPL): {freq_vs_vf[int(0.01/0.03*15)]:.2f} rad/s")
    print(f"Enhancement factor: {freq_vs_vf[int(0.01/0.03*15)]/freq_vs_vf[0]:.2f}")
    print(f"Best distribution pattern: {max(pattern_freqs.items(), key=lambda x: x[1][0])[0]}")
    print(f"SSSS vs CCCC frequency ratio: {bc_freqs['CCCC'][0]/bc_freqs['SSSS'][0]:.2f}")

def interactive_analysis():
    """Interactive analysis tool for custom configurations."""
    print("\nInteractive Analysis Tool")
    print("=" * 30)
    
    epoxy, graphene, gpl_geometry = create_material_properties()
    grc = GrapheneReinforcedComposite(epoxy, graphene, gpl_geometry)
    
    # Get user input
    print("\nPlate Configuration:")
    length = float(input("Plate length (m) [0.1]: ") or 0.1)
    width = float(input("Plate width (m) [0.1]: ") or 0.1)
    thickness = float(input("Plate thickness (m) [0.002]: ") or 0.002)
    
    print("\nGRC Configuration:")
    vf = float(input("GPL volume fraction (0-0.05) [0.01]: ") or 0.01)
    
    pattern_map = {'1': DistributionPattern.UD, '2': DistributionPattern.FG_V,
                   '3': DistributionPattern.FG_X, '4': DistributionPattern.FG_O}
    print("Distribution patterns: 1=UD, 2=FG-V, 3=FG-X, 4=FG-O")
    pattern_choice = input("Choose pattern [3]: ") or '3'
    pattern = pattern_map[pattern_choice]
    
    print("\nBoundary Conditions:")
    bc_map = {'1': BoundaryCondition.SIMPLY_SUPPORTED, '2': BoundaryCondition.CLAMPED}
    print("1=Simply Supported, 2=Clamped")
    bc_choice = input("Choose boundary condition [1]: ") or '1'
    bc = bc_map[bc_choice]
    
    # Create and analyze plate
    layers = [Layer(thickness=thickness, material=epoxy, is_grc=True,
                   grc_material=grc, grc_pattern=pattern,
                   grc_volume_fraction=vf)]
    
    plate = LaminatedPlate(layers, length=length, width=width)
    solver = PlateVibrationSolver(plate, n_points_x=9, n_points_y=9)
    
    bc_config = ((bc, bc), (bc, bc))
    freqs, modes = solver.solve_eigenvalue_problem(bc_config[0], bc_config[1], n_modes=5)
    
    print(f"\nResults:")
    print(f"Natural Frequencies (rad/s):")
    for i, freq in enumerate(freqs):
        dim_freq = solver.dimensionless_frequency(freq, epoxy)
        print(f"  Mode {i+1}: {freq:.2f} rad/s (Ω = {dim_freq:.4f})")
    
    # Compare with pure epoxy
    layers_epoxy = [Layer(thickness=thickness, material=epoxy)]
    plate_epoxy = LaminatedPlate(layers_epoxy, length=length, width=width)
    solver_epoxy = PlateVibrationSolver(plate_epoxy, n_points_x=9, n_points_y=9)
    freqs_epoxy, _ = solver_epoxy.solve_eigenvalue_problem(bc_config[0], bc_config[1], n_modes=1)
    
    enhancement = freqs[0] / freqs_epoxy[0]
    print(f"\nEnhancement factor: {enhancement:.2f}")
    print(f"Effective modulus: {grc.effective_modulus(vf):.2f} GPa")

def main():
    """Main function with command line interface."""
    parser = argparse.ArgumentParser(
        description="Dynamic Analysis of FG-GRC Laminated Plates",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python main.py --validation          # Run validation studies
  python main.py --parametric          # Run parametric studies  
  python main.py --interactive         # Interactive analysis
  python main.py --all                 # Run all analyses
        """
    )
    
    parser.add_argument('--validation', action='store_true',
                       help='Run validation studies against paper results')
    parser.add_argument('--parametric', action='store_true', 
                       help='Run comprehensive parametric studies')
    parser.add_argument('--interactive', action='store_true',
                       help='Run interactive analysis tool')
    parser.add_argument('--all', action='store_true',
                       help='Run all analyses')
    parser.add_argument('--output-dir', type=str, default='.',
                       help='Output directory for results')
    
    args = parser.parse_args()
    
    if not any([args.validation, args.parametric, args.interactive, args.all]):
        parser.print_help()
        return
    
    output_dir = Path(args.output_dir)
    output_dir.mkdir(exist_ok=True)
    
    print("Dynamic Analysis of Composite Laminated Plates")
    print("with Functionally Graded Graphene-Reinforced Layer")
    print("Based on Wang et al. (2015)")
    print("=" * 60)
    
    if args.validation or args.all:
        print("\nRunning Validation Studies...")
        validation = ValidationStudy()
        validation.isotropic_plate_validation()
        validation.gpl_weight_fraction_study() 
        validation.distribution_pattern_study()
        validation.boundary_condition_study()
        create_validation_plots(validation)
    
    if args.parametric or args.all:
        print("\nRunning Parametric Studies...")
        parametric_study(output_dir)
    
    if args.interactive or args.all:
        interactive_analysis()
    
    print("\nAnalysis completed successfully!")
    print(f"Results saved to: {output_dir.absolute()}")

if __name__ == "__main__":
    main()